import torch
import numpy as np


def station_tp_mae_loss(input, target):
    # aa = input
    # print(aa)
    n = torch.abs(input - target)
    ret = torch.where(target > 0.9, n * 2, n)
    ret = torch.mean(ret)
    return ret


def ts_station(x, y, idx=0, threshold=0):
    assert x.shape == y.shape  # x预报，y实况
    ths = [0.1, 1, 5, 10, 15, 20]

    # scores = []  # TS评分
    scores = np.zeros(len(ths))
    for i, th in enumerate(ths):
        # a = np.sum((x >= th)[idx] & (y >= th)[idx])  # 预报正确 [..., idx[0], idx[1]]
        a = np.sum(((x >= th) & (y >= th))[idx])  # 预报正确
        b = np.sum(((x >= th) & (y < th))[idx])  # 空报
        c = np.sum(((x < th) & (y >= th))[idx])  # 漏报
        d = np.sum(((x < th) & (y < th))[idx])  # 无降水预报正确

        # all = a + b + c + d
        # assert (x.size == all.sum())

        if a + b + c > 0:
            _ts = a / (a + b + c)
            # print(_ts)
            scores[i] = _ts
    return scores


def weighted_mae_loss(input, target):
    n = torch.abs(input - target)
    torch.where(input > 1, n * input * 100, n * 10)
    # torch.where(input > 4, n * input * 2000, n)
    # torch.where(input < 0.8, n , n)
    ret = torch.mean(n)
    return ret


def weighted_mae_loss2(input, target):
    n = torch.abs(input - target)
    torch.where(input > 1, n * input * 10, n)
    torch.where(input > 5, n * input * 10, n)
    # torch.where(input > 4, n * input * 2000, n)
    # torch.where(input < 0.8, n , n)
    ret = torch.mean(n)
    return ret


def area_loss(input, target):
    input_area = torch.sign(input)
    target_area = torch.sign(target)
    ret = torch.abs(input_area - target_area).mean() + 1
    ret = torch.where(ret > 0, torch.log(ret), ret * 0)
    # ret = torch.clamp(ret, min=0.0, max=1.0)
    return ret


def iou_loss(input, target):
    union = input + target
    inter = input * target

    ious = 1 - torch.sum(inter) / torch.sum(union)
    ious = torch.clamp(ious, min=0.0, max=1.0)
    return ious


def grid_cell_regularizer(input, target):
    """Grid cell regularizer."""
    # m = 24 / 5
    # _w = target + 1
    # _w[_w < m] = m
    loss = torch.mean(torch.abs((input - target)))
    return loss
